/*********************IMPORTANT DRAKVUF LICENSE TERMS***********************
 *                                                                         *
 * DRAKVUF (C) 2014-2022 Tamas K Lengyel.                                  *
 * Tamas K Lengyel is hereinafter referred to as the author.               *
 * This program is free software; you may redistribute and/or modify it    *
 * under the terms of the GNU General Public License as published by the   *
 * Free Software Foundation; Version 2 ("GPL"), BUT ONLY WITH ALL OF THE   *
 * CLARIFICATIONS AND EXCEPTIONS DESCRIBED HEREIN.  This guarantees your   *
 * right to use, modify, and redistribute this software under certain      *
 * conditions.  If you wish to embed DRAKVUF technology into proprietary   *
 * software, alternative licenses can be acquired from the author.         *
 *                                                                         *
 * Note that the GPL places important restrictions on "derivative works",  *
 * yet it does not provide a detailed definition of that term.  To avoid   *
 * misunderstandings, we interpret that term as broadly as copyright law   *
 * allows.  For example, we consider an application to constitute a        *
 * derivative work for the purpose of this license if it does any of the   *
 * following with any software or content covered by this license          *
 * ("Covered Software"):                                                   *
 *                                                                         *
 * o Integrates source code from Covered Software.                         *
 *                                                                         *
 * o Reads or includes copyrighted data files.                             *
 *                                                                         *
 * o Is designed specifically to execute Covered Software and parse the    *
 * results (as opposed to typical shell or execution-menu apps, which will *
 * execute anything you tell them to).                                     *
 *                                                                         *
 * o Includes Covered Software in a proprietary executable installer.  The *
 * installers produced by InstallShield are an example of this.  Including *
 * DRAKVUF with other software in compressed or archival form does not     *
 * trigger this provision, provided appropriate open source decompression  *
 * or de-archiving software is widely available for no charge.  For the    *
 * purposes of this license, an installer is considered to include Covered *
 * Software even if it actually retrieves a copy of Covered Software from  *
 * another source during runtime (such as by downloading it from the       *
 * Internet).                                                              *
 *                                                                         *
 * o Links (statically or dynamically) to a library which does any of the  *
 * above.                                                                  *
 *                                                                         *
 * o Executes a helper program, module, or script to do any of the above.  *
 *                                                                         *
 * This list is not exclusive, but is meant to clarify our interpretation  *
 * of derived works with some common examples.  Other people may interpret *
 * the plain GPL differently, so we consider this a special exception to   *
 * the GPL that we apply to Covered Software.  Works which meet any of     *
 * these conditions must conform to all of the terms of this license,      *
 * particularly including the GPL Section 3 requirements of providing      *
 * source code and allowing free redistribution of the work as a whole.    *
 *                                                                         *
 * Any redistribution of Covered Software, including any derived works,    *
 * must obey and carry forward all of the terms of this license, including *
 * obeying all GPL rules and restrictions.  For example, source code of    *
 * the whole work must be provided and free redistribution must be         *
 * allowed.  All GPL references to "this License", are to be treated as    *
 * including the terms and conditions of this license text as well.        *
 *                                                                         *
 * Because this license imposes special exceptions to the GPL, Covered     *
 * Work may not be combined (even as part of a larger work) with plain GPL *
 * software.  The terms, conditions, and exceptions of this license must   *
 * be included as well.  This license is incompatible with some other open *
 * source licenses as well.  In some cases we can relicense portions of    *
 * DRAKVUF or grant special permissions to use it in other open source     *
 * software.  Please contact tamas.k.lengyel@gmail.com with any such       *
 * requests.  Similarly, we don't incorporate incompatible open source     *
 * software into Covered Software without special permission from the      *
 * copyright holders.                                                      *
 *                                                                         *
 * If you have any questions about the licensing restrictions on using     *
 * DRAKVUF in other works, are happy to help.  As mentioned above,         *
 * alternative license can be requested from the author to integrate       *
 * DRAKVUF into proprietary applications and appliances.  Please email     *
 * tamas.k.lengyel@gmail.com for further information.                      *
 *                                                                         *
 * If you have received a written license agreement or contract for        *
 * Covered Software stating terms other than these, you may choose to use  *
 * and redistribute Covered Software under those terms instead of these.   *
 *                                                                         *
 * Source is provided to this software because we believe users have a     *
 * right to know exactly what a program is going to do before they run it. *
 * This also allows you to audit the software for security holes.          *
 *                                                                         *
 * Source code also allows you to port DRAKVUF to new platforms, fix bugs, *
 * and add new features.  You are highly encouraged to submit your changes *
 * on https://github.com/tklengyel/drakvuf, or by other methods.           *
 * By sending these changes, it is understood (unless you specify          *
 * otherwise) that you are offering unlimited, non-exclusive right to      *
 * reuse, modify, and relicense the code.  DRAKVUF will always be          *
 * available Open Source, but this is important because the inability to   *
 * relicense code has caused devastating problems for other Free Software  *
 * projects (such as KDE and NASM).                                        *
 * To specify special license conditions of your contributions, just say   *
 * so when you send them.                                                  *
 *                                                                         *
 * This program is distributed in the hope that it will be useful, but     *
 * WITHOUT ANY WARRANTY; without even the implied warranty of              *
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the DRAKVUF   *
 * license file for more details (it's in a COPYING file included with     *
 * DRAKVUF, and also available from                                        *
 * https://github.com/tklengyel/drakvuf/COPYING)                           *
 *                                                                         *
 ***************************************************************************/

#include <glib.h>
#include <libvmi/libvmi.h>

#include "plugins/output_format.h"
#include "exploitmon.h"
#include "private.h"

static bool read_kernel_addr(drakvuf_t drakvuf, addr_t read_addr, addr_t* out_addr)
{
    vmi_lock_guard vmi(drakvuf);
    ACCESS_CONTEXT(ctx);
    ctx.translate_mechanism = VMI_TM_PROCESS_DTB;
    ctx.addr = read_addr;

    if (VMI_SUCCESS != vmi_pid_to_dtb(vmi, 4, &ctx.dtb))
        return false;
    if (VMI_SUCCESS != vmi_read_addr(vmi, &ctx, out_addr))
        return false;
    return true;
}

static void process_visitor(drakvuf_t drakvuf, addr_t process, void* ctx)
{
    auto plugin = static_cast<exploitmon*>(ctx);
    // Read process token, pid
    addr_t token, pid;
    if (!read_kernel_addr(drakvuf, process + plugin->offsets[EPROCESS_TOKEN], &token))
    {
        PRINT_DEBUG("[EXPLOITMON] Failed to read process token\n");
        throw -1;
    }
    if (!read_kernel_addr(drakvuf, process + plugin->offsets[EPROCESS_UNIQUE_PROCESS_ID], &pid))
    {
        PRINT_DEBUG("[EXPLOITMON] Failed to read process pid\n");
        throw -1;
    }
    // Strip ref count bits
    token = token & plugin->ex_fast_ref_mask;
    plugin->live_processes.push_back(std::make_pair(pid, token));
}

static void collect_eprocess(drakvuf_t drakvuf, addr_t process, void* ctx)
{
    auto data = static_cast<std::vector<addr_t>*>(ctx);
    data->push_back(process);
}

/// Hook physical page @addr for @access_type
static bool hook_page(
    exploitmon* plugin,
    drakvuf_t drakvuf,
    addr_t addr,
    vmi_mem_access_t access_type,
    event_response_t (*callback)(drakvuf_t, drakvuf_trap_info_t*),
    addr_t data = 0)
{
    auto trap = new (std::nothrow) drakvuf_trap_t();
    if (!trap)
        return false;

    trap->type = MEMACCESS;
    trap->memaccess.gfn = addr >> 12;
    trap->memaccess.type = PRE;
    trap->memaccess.access = access_type;
    trap->cb = callback;
    plugin->register_trap<page_fault_call_result>(trap);
    auto params = get_trap_params<page_fault_call_result>(trap);
    params->addr = data;
    return true;
}

static bool translate_ksym2p(vmi_instance_t vmi, const char* symbol, addr_t* addr)
{
    addr_t temp_va;
    if (VMI_SUCCESS != vmi_translate_ksym2v(vmi, symbol, &temp_va))
    {
        PRINT_DEBUG("[EXPLOITMON] Failed to translate symbol to virtual address\n");
        return false;
    }
    if (VMI_SUCCESS != vmi_translate_kv2p(vmi, temp_va, addr))
    {
        PRINT_DEBUG("[EXPLOITMON] Failed to translate virtual address to physical\n");
        return false;
    }
    return true;
}

static event_response_t terminate_process_hook_cb(drakvuf_t drakvuf, drakvuf_trap_info_t* info)
{
    auto plugin = get_trap_plugin<exploitmon>(info);
    // HANDLE ProcessHandle
    uint64_t process_handle = drakvuf_get_function_argument(drakvuf, info, 1);

    if (process_handle != ~0ULL)
    {
        PRINT_DEBUG("[EXPLOITMON] Process handle not pointing to self, ignore\n");
        return VMI_EVENT_RESPONSE_NONE;
    }

    addr_t process = info->attached_proc_data.base_addr;

    addr_t token, pid;
    if (!read_kernel_addr(drakvuf, process + plugin->offsets[EPROCESS_TOKEN], &token))
    {
        PRINT_DEBUG("[EXPLOITMON] Failed to read process token\n");
        return VMI_EVENT_RESPONSE_NONE;
    }
    if (!read_kernel_addr(drakvuf, process + plugin->offsets[EPROCESS_UNIQUE_PROCESS_ID], &pid))
    {
        PRINT_DEBUG("[EXPLOITMON] Failed to read process pid\n");
        return VMI_EVENT_RESPONSE_NONE;
    }
    // Strip refcount bits
    token = token & plugin->ex_fast_ref_mask;
    // Compare token on process termination with token at process creation
    // Perhaps there is a better way..
    int i = 0;
    for (const auto& [p_pid, p_token] : plugin->live_processes)
    {
        if (p_pid == pid)
        {
            if (p_token != token)
            {
                fmt::print(plugin->format, "exploitmon", drakvuf, info,
                    keyval("Reason", fmt::Qstr("Token modification detected")));
            }
            // Don't forget to remove it!
            plugin->live_processes.erase(plugin->live_processes.begin() + i);
            break;
        }
        i++;
    }
    return VMI_EVENT_RESPONSE_NONE;
}

static event_response_t insert_process_hook_cb(drakvuf_t drakvuf, drakvuf_trap_info_t* info)
{
    addr_t process = drakvuf_get_function_argument(drakvuf, info, 1);
    auto plugin = get_trap_plugin<exploitmon>(info);
    // Read process token and pid
    addr_t token, pid;
    if (!read_kernel_addr(drakvuf, process + plugin->offsets[EPROCESS_TOKEN], &token))
    {
        PRINT_DEBUG("[EXPLOITMON] Failed to read process token\n");
        throw -1;
    }
    if (!read_kernel_addr(drakvuf, process + plugin->offsets[EPROCESS_UNIQUE_PROCESS_ID], &pid))
    {
        PRINT_DEBUG("[EXPLOITMON] Failed to read process pid\n");
        throw -1;
    }
    // Strip ref count bits
    token = token & plugin->ex_fast_ref_mask;
    plugin->live_processes.push_back(std::make_pair(pid, token));
    return VMI_EVENT_RESPONSE_NONE;
}

static event_response_t final_token_check_cb(drakvuf_t drakvuf, drakvuf_trap_info_t* info)
{
    auto plugin = get_trap_plugin<exploitmon>(info);
    if (plugin->is_stopping() && !plugin->done_final_analysis)
    {
        PRINT_DEBUG("[EXPLOITMON] Making final analysis\n");
        // Copy past processes
        auto past_processes = plugin->live_processes;
        plugin->live_processes.clear();
        // Enumerate existing processes
        drakvuf_enumerate_processes(drakvuf, process_visitor, static_cast<void*>(plugin));
        // Collect eprocess
        auto process_list = new std::vector<addr_t>;
        drakvuf_enumerate_processes(drakvuf, collect_eprocess, static_cast<void*>(process_list));
        // Compare
        for (const auto& [pid, token] : plugin->live_processes)
        {
            for (const auto& [p_pid, p_token] : past_processes)
            {
                // If token mismatch
                if (p_pid == pid && token != p_token)
                {
                    // Get corresponding process
                    for (const auto& process : *process_list)
                    {
                        addr_t temp_pid = 0;
                        if (!read_kernel_addr(drakvuf, process + plugin->offsets[EPROCESS_UNIQUE_PROCESS_ID], &temp_pid))
                        {
                            PRINT_DEBUG("[EXPLOITMON] Failed to read process token\n");
                            continue;
                        }
                        if (temp_pid == pid)
                        {
                            // Save old proc_data
                            auto temp_attach_data = info->attached_proc_data;
                            auto temp_data = info->proc_data;

                            // Fill new proc_data
                            proc_data_t data = {};
                            if (!drakvuf_get_process_data(drakvuf, process, &data))
                            {
                                PRINT_DEBUG("[EXPLOITMON] Failed to get process data\n");
                                break;
                            }
                            // Replace info proc_data
                            info->attached_proc_data = data;
                            info->proc_data = data;

                            fmt::print(plugin->format, "exploitmon", drakvuf, info,
                                keyval("Reason", fmt::Qstr("Token modification detected")));

                            // Dont forget to free process name!
                            g_free(const_cast<char*>(data.name));

                            // Replace back old data
                            info->attached_proc_data = temp_attach_data;
                            info->proc_data = temp_data;

                            break;
                        }
                    }
                }
            }
        }

        delete process_list;

        plugin->done_final_analysis = true;
    }
    return VMI_EVENT_RESPONSE_NONE;
}

static event_response_t execute_faulted_cb(drakvuf_t drakvuf, drakvuf_trap_info_t* info)
{
    auto plugin = get_trap_plugin<exploitmon>(info);
    auto params = get_trap_params<page_fault_call_result>(info->trap);
    auto vaddr = params->addr;

    plugin->destroy_trap(info->trap);

    page_info_t p_info = {};
    {
        vmi_lock_guard vmi(drakvuf);
        if (VMI_SUCCESS != vmi_pagetable_lookup_extended(vmi, info->regs->cr3, vaddr, &p_info))
            return VMI_EVENT_RESPONSE_NONE;
    }
    auto page_user  = (p_info.x86_ia32e.pte_value) & 4;
    privilege_mode_t mode;
    if (!drakvuf_get_current_thread_previous_mode(drakvuf, info, &mode))
        return VMI_EVENT_RESPONSE_NONE;

    if (page_user ^ mode)
    {
        fmt::print(plugin->format, "exploitmon", drakvuf, info,
            keyval("Reason", fmt::Qstr("Kernel thread execution")));

        PRINT_DEBUG("[EXPLOITMON] Kernel thread executing user memory\n");
    }
    return VMI_EVENT_RESPONSE_NONE;
}

static event_response_t page_fault_return_hook_cb(drakvuf_t drakvuf, drakvuf_trap_info_t* info)
{
    auto params = get_trap_params<page_fault_call_result>(info);
    if (!params->verify_result_call_params(drakvuf, info))
        return VMI_EVENT_RESPONSE_NONE;
    drakvuf_remove_trap(drakvuf, info->trap, nullptr);
    page_info_t page = {};
    {
        vmi_lock_guard vmi(drakvuf);
        if (VMI_SUCCESS != vmi_pagetable_lookup_extended(vmi, info->regs->cr3, params->addr, &page))
            return VMI_EVENT_RESPONSE_NONE;
    }

    if (!hook_page(get_trap_plugin<exploitmon>(info), drakvuf, page.paddr, VMI_MEMACCESS_X, execute_faulted_cb, params->addr))
    {
        PRINT_DEBUG("[EXPLOITMON] Failed to hook page in MmAccessFault\n");
    }
    return VMI_EVENT_RESPONSE_NONE;
}

static event_response_t page_fault_hook_cb(drakvuf_t drakvuf, drakvuf_trap_info_t* info)
{
    addr_t vaddr = drakvuf_get_function_argument(drakvuf, info, 2);
    // Upper bits -> kernel memory (not interested)
    if (vaddr & (1UL << 63))
        return VMI_EVENT_RESPONSE_NONE;

    auto plugin = get_trap_plugin<exploitmon>(info);
    auto trap = plugin->register_trap<page_fault_call_result>(
            info,
            page_fault_return_hook_cb,
            breakpoint_by_pid_searcher());

    if (!trap)
        return VMI_EVENT_RESPONSE_NONE;

    auto params = get_trap_params<page_fault_call_result>(trap);
    params->set_result_call_params(info);
    params->addr = vaddr;
    return VMI_EVENT_RESPONSE_NONE;
}

static event_response_t hal_table_overwrite_cb(drakvuf_t drakvuf, drakvuf_trap_info_t* info)
{
    PRINT_DEBUG("[EXPLOITMON] In hal overwrite\n");
    auto plugin = get_trap_plugin<exploitmon>(info);
    auto params = get_trap_params<page_fault_call_result>(info->trap);
    addr_t hal_table_pa = params->addr;
    // https://www.geoffchappell.com/studies/windows/km/ntoskrnl/inc/ntos/hal/hal_dispatch.htm
    if (info->trap_pa >= hal_table_pa && info->trap_pa < hal_table_pa + 0x100)
    {
        fmt::print(plugin->format, "exploitmon", drakvuf, info,
            keyval("Reason", fmt::Qstr("HalDispatchTable overwrite detected")));
    }
    return VMI_EVENT_RESPONSE_NONE;
}

/// Enumerate existing processes
/// Hook HalDispatchTable
/// Hook PspTerminateProcess
/// Hook PspInsertProcess
exploitmon::exploitmon(drakvuf_t drakvuf, struct exploitmon_config* c, output_format_t output)
    : pluginex(drakvuf, output), format(output), offsets(new size_t[__OFFSET_MAX]), done_final_analysis(false)
{
    bool is32bit = (drakvuf_get_page_mode(drakvuf) != VMI_PM_IA32E);
    ex_fast_ref_mask = is32bit ? ~7 : ~0xf;

    // Collect kernel struct member offsets
    if (!drakvuf_get_kernel_struct_members_array_rva(drakvuf, offset_names, __OFFSET_MAX, offsets))
    {
        PRINT_DEBUG("[EXPLOITMON] Failed to get kernel struct member offsets\n");
        throw -1;
    }
    // Enumerate existing processes
    drakvuf_enumerate_processes(drakvuf, process_visitor, static_cast<void*>(this));

    addr_t hal_table_pa;
    {
        vmi_lock_guard vmi(drakvuf);

        if (!translate_ksym2p(vmi, "HalDispatchTable", &hal_table_pa))
        {
            PRINT_DEBUG("[EXPLOITMON] Failed to translate HalDispatchTable to physical address\n");
            throw -1;
        }
    }
    // Hook mem write on HalDispatchTable page
    if (!hook_page(this, drakvuf, hal_table_pa, VMI_MEMACCESS_W, hal_table_overwrite_cb, hal_table_pa))
    {
        PRINT_DEBUG("[EXPLOITMON] Failed to hook HalDispatchTable\n");
        throw -1;
    }

    breakpoint_in_system_process_searcher bp;
    if (!register_trap(nullptr, terminate_process_hook_cb, bp.for_syscall_name("NtTerminateProcess")))
    {
        PRINT_DEBUG("[EXPLOITMON] Failed to hook NtTerminateProcess\n");
        throw -1;
    }
    if (!register_trap(nullptr, insert_process_hook_cb, bp.for_syscall_name("PspInsertProcess")))
    {
        PRINT_DEBUG("[EXPLOITMON] Faild to hook PspInsertProcess\n");
        throw -1;
    }
    if (c->enable_k2u)
    {
        if (!register_trap(nullptr, page_fault_hook_cb, bp.for_syscall_name("MmAccessFault")))
        {
            PRINT_DEBUG("[EXPLOITMON] Faild to hook MmAccessFault\n");
            throw -1;
        }
    }
}

exploitmon::~exploitmon()
{
    delete[] offsets;
    destroy_all_traps();
}

bool exploitmon::stop_impl()
{
    if (!is_stopping() && !done_final_analysis)
    {
        // Hook dummy function so we can make final analysis
        breakpoint_in_system_process_searcher bp;
        if (!register_trap(nullptr, final_token_check_cb, bp.for_syscall_name("NtClose")))
        {
            // Skip final analysis
            PRINT_DEBUG("[EXPLOITMON] Failed to hook NtClose\n");
            done_final_analysis = true;
            return pluginex::stop_impl();
        }
        // Return status `Pending`
        return false;
    }
    if (done_final_analysis)
        return pluginex::stop_impl();
    return false;
}
